import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_wine
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

class MyRandomForest:
    def __init__(self, n_trees=10, max_depth=None, min_samples_split=2):
        self.n_trees = n_trees
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.trees = []

    def bootstrap_sample(self, X, y):
        n_samples = X.shape[0]
        idxs = np.random.choice(n_samples, size=n_samples, replace=True)
        return X[idxs], y[idxs]

    def fit(self, X, y):
        self.trees = []
        for _ in range(self.n_trees):
            tree = DecisionTreeClassifier(
                max_depth=self.max_depth,
                min_samples_split=self.min_samples_split
            )
            X_sample, y_sample = self.bootstrap_sample(X, y)
            tree.fit(X_sample, y_sample)
            self.trees.append(tree)

    def predict(self, X):
        tree_preds = np.array([tree.predict(X) for tree in self.trees])
        return np.array([np.bincount(tree_preds[:, i]).argmax() 
                        for i in range(len(X))])

def load_and_prepare_data():
    """加载wine数据集并准备数据"""
    wine = load_wine()
    X = wine.data
    y = wine.target
    return train_test_split(X, y, test_size=0.2, random_state=42)

def train_and_evaluate_models(X_train, X_test, y_train, y_test):
    """训练和评估三种模型"""
    # 训练决策树
    dt = DecisionTreeClassifier(random_state=42)
    dt.fit(X_train, y_train)
    dt_pred = dt.predict(X_test)
    dt_accuracy = accuracy_score(y_test, dt_pred)

    # 训练sklearn随机森林
    rf = RandomForestClassifier(n_estimators=10, random_state=42)
    rf.fit(X_train, y_train)
    rf_pred = rf.predict(X_test)
    rf_accuracy = accuracy_score(y_test, rf_pred)

    # 训练自定义随机森林
    my_rf = MyRandomForest(n_trees=10)
    my_rf.fit(X_train, y_train)
    my_rf_pred = my_rf.predict(X_test)
    my_rf_accuracy = accuracy_score(y_test, my_rf_pred)

    return {
        'Decision Tree': dt_accuracy,
        'Sklearn Random Forest': rf_accuracy,
        'My Random Forest': my_rf_accuracy
    }

def plot_results(accuracies):
    """可视化三种模型的准确率比较"""
    plt.figure(figsize=(10, 6))
    models = list(accuracies.keys())
    accs = list(accuracies.values())
    
    bars = plt.bar(models, accs)
    plt.title('模型准确率比较')
    plt.ylabel('准确率')
    
    # 为每个柱状图添加具体数值
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height,
                 f'{height:.4f}',
                 ha='center', va='bottom')
    
    plt.xticks(rotation=15)
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 加载数据集
    X_train, X_test, y_train, y_test = load_and_prepare_data()
    
    # 训练和评估模型
    accuracies = train_and_evaluate_models(X_train, X_test, y_train, y_test)
    
    # 打印结果
    print("\n模型准确率比较:")
    for model, accuracy in accuracies.items():
        print(f"{model}: {accuracy:.4f}")
    
    # 可视化结果
    plot_results(accuracies)